import numpy as np
import pandas as pd
from scipy.optimize import minimize, basinhopping

class LPPLSModel:
    def __init__(self, data):
        self.data = data
        self.tc = None
        self.m = None
        self.omega = None
        self.phi = None
        self.A = None
        self.B = None
        self.C1 = None
        self.C2 = None

    def lppls(self, t, tc, m, omega, phi, A, B, C1, C2):
        return A + B * ((tc - t) ** m) * (1 + C1 * np.cos(omega * np.log(tc - t) + phi)) + C2 * (tc - t)

    def objective(self, params, t, price):
        tc, m, omega, phi, A, B, C1, C2 = params
        return np.sum((price - self.lppls(t, tc, m, omega, phi, A, B, C1, C2)) ** 2)

    def fit(self, initial_guess):
        t = np.arange(len(self.data))
        price = self.data.values
        result = minimize(self.objective, initial_guess, args=(t, price),
                         bounds=[(len(self.data)*0.9, len(self.data)*1.1), (0, 1), (0, 20), (0, 2*np.pi),
                                 (None, None), (None, None), (None, None), (None, None)])
        self.tc, self.m, self.omega, self.phi, self.A, self.B, self.C1, self.C2 = result.x
        return result.x

    def predict(self, days):
        t = np.arange(len(self.data), len(self.data) + days)
        return self.lppls(t, self.tc, self.m, self.omega, self.phi, self.A, self.B, self.C1, self.C2)

    def auto_optimize(self, initial_guess, niter=100, T=1.0, stepsize=0.5, threshold=1e-5):
        """
        自动优化参数
        :param initial_guess: 初始参数猜测
        :param niter: 迭代次数
        :param T: 温度参数
        :param stepsize: 步长
        :param threshold: 收敛阈值
        """
        t = np.arange(len(self.data))
        price = self.data.values
        
        def callback(x, f, accept):
            if f < threshold:
                return True

        minimizer_kwargs = {
            'args': (t, price),
            'bounds': [(len(self.data)*0.9, len(self.data)*1.1), (0, 1), (0, 20), (0, 2*np.pi),
                      (None, None), (None, None), (None, None), (None, None)]
        }
        
        result = basinhopping(self.objective, initial_guess, niter=niter, T=T,
                             stepsize=stepsize, minimizer_kwargs=minimizer_kwargs,
                             callback=callback)
        self.tc, self.m, self.omega, self.phi, self.A, self.B, self.C1, self.C2 = result.x
        return result.x

    def plot(self, days=30):
        import matplotlib.pyplot as plt
        # 历史拟合数据
        t_history = np.arange(len(self.data))
        fitted_values = self.lppls(t_history, self.tc, self.m, self.omega, self.phi, self.A, self.B, self.C1, self.C2)
        # 未来预测数据
        t_future = np.arange(len(self.data), len(self.data) + days)
        predicted_values = self.predict(days)
        # 绘图
        fig, (ax1, ax2) = plt.subplots(nrows=2, ncols=1, sharex=True, figsize=(18, 10))
        # 绘制原始数据和拟合曲线
        ax1.plot(t_history, self.data.values, 'b-', label='实际数据')
        ax1.plot(t_history, fitted_values, 'r--', label='拟合曲线')
        ax1.plot(t_future, predicted_values, 'g--', label='预测曲线')
        ax1.axvline(x=len(self.data), color='gray', linestyle=':', label='预测起点')
        ax1.set_xlabel('时间')
        ax1.set_ylabel('价格')
        ax1.set_title('LPPLS模型拟合与预测')
        ax1.legend()
        ax1.grid(True)
        # 绘制bubble indicator
        ax2_0 = ax2.twinx()
        ax2.plot(t_history, self.data.values, 'b-', linewidth=0.75)
        ax2_0.plot(t_history, fitted_values, 'r--', label='bubble indicator', alpha=0.5)
        ax2.set_xlabel('时间')
        ax2.set_ylabel('价格')
        ax2_0.set_ylabel('bubble indicator')
        ax2_0.legend(loc=2)
        ax2.grid(True)
        plt.xticks(rotation=45)
        plt.savefig('lppls_bubble_plot.png')
        plt.show()

if __name__ == '__main__':
    # 加载黄金价格数据
    data = pd.read_excel('伦敦市场黄金现货价.xlsx')
    data = data.iloc[-1000:, :]
    # 假设价格在'Price'列
    model = LPPLSModel(data['Price'])
    # 初始参数猜测
    initial_guess = [len(data)*1.05, 0.5, 10, 1, data['Price'].mean(), -1, 0.1, 0.1]
    # 拟合模型
    # 使用自动优化方法
    optimized_params = model.auto_optimize(initial_guess)
    print('优化后的参数:', optimized_params)
    # 使用优化后的参数进行预测
    model.fit(optimized_params)
    model.plot()
    # 预测未来30天
    predictions = model.predict(30)
    print(predictions)